""" Configurable. Run once then poison channels. """
from common import *

import pylab

default_data = 1

def setup(data = default_data):
    import wx
    dlg = wx.TextEntryDialog(
        None, 'Enter number of cannons',
        'Configure EnergyCollector Process', str(data))
    if dlg.ShowModal() == wx.ID_OK:
        try:
            data = eval(dlg.GetValue())
        except:
            data = dlg.GetValue()
    dlg.Destroy()
    return data

def EnergyCollectorFunc(energyIn, imageIn , data = default_data):

    image = imageIn()
    xsize = len(image[0])
    ysize = len(image)

    prev_image = [(xsize)*[0] for i in range(ysize)]
    for y in range(ysize):
        for x in range(xsize):
            prev_image[y][x] = (float(image[y][x])/255.0,float(image[y][x])/255.0,float(image[y][x])/255.0,1)

    pylab.imshow(image, cmap = pylab.cm.gray)
    pylab.show() 

    while True:
        cannonCount = data
                
        allEnergy = [(xsize)*[0] for i in range(ysize)]

        # Process cannons
        for i in range(cannonCount):
            energy = energyIn()

            # Add to all energy
            for y in range(ysize):
                for x in range(xsize):
                    allEnergy[y][x] = allEnergy[y][x] + energy[y][x]

        # Draw result
        
        scale=[10, 20, 30, 40, 50, 60, 70, 80, 90]
        #cmap  =[(0,0,85,1), (0,0,170,1), (0,0,255,1), (0,85,0,1), (0,170,0,1), (0,255,0,1),  (85,0,0,1),  (170,0,0,1), (255,0,0,1)]
        cmap  =[(0,0,85.0/255.0,1), (0,0,170.0/255.0,1), (0,0,255.0/255.0,1), (0,85.0/255.0,0,1), (0,170.0/255.0,0,1), (0,255.0/255.0,0,1),  (85.0/255.0,0,0,1),  (170.0/255.0,0,0,1), (255.0/255.0,0,0,1)]
        def fpos(x,scale):
            for i in range(len(scale)):
                if x<=scale[i]: return i
            return i


        out_img = [(xsize)*[0] for i in range(ysize)]
        for y in range(ysize):
            for x in range(xsize):
                if (allEnergy[y][x]>0):
                    out_img[y][x] = cmap[fpos(allEnergy[y][x], scale)]
                else:
                    out_img[y][x] = (float(image[y][x])/255.0,float(image[y][x])/255.0,float(image[y][x])/255.0,1)
        
        pylab.imshow(pylab.concatenate([prev_image, out_img], 1))
        pylab.show()

        prev_image = out_img    

